import json
from operator import itemgetter
from typing import Any, Dict, List, Optional, TypedDict

from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool, render_text_description
from langchain_core.runnables import RunnableConfig
from langchain_core.utils.function_calling import convert_to_openai_tool

from ChatGLM_new import tongyi_llm,zhipu_llm,xinghuo_llm,local_llm,sspy_llm,ollama_llm,tongyi_llm_openai,kimi_llm,kimi_llm_openai
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.agents import AgentExecutor, create_tool_calling_agent

class RoadModel(BaseModel):
    """Information about a road."""
    roadName: Optional[str] = Field(default=None, description="道路名称")
    roadSide: Optional[str] = Field(default=None, description="道路路侧")
    dimmingValue: Optional[int] = Field(default=None, description="调光值,开灯时默认为100,关灯时默认为0")

class LamppostModel(BaseModel):
    """Information about pole."""
    lamppostCode: Optional[str] = Field(default=None, description="灯杆编号")
    dimmingValue: Optional[int] = Field(default=None, description="调光值,开灯时默认为100,关灯时默认为0")

@tool
def road_promote(roadModelList: List[RoadModel]):
    """ 按道路名称进行开关灯。参数格式为json:{"roadModelList": [{"roadName": "舜华路", "roadSide": "", "dimmingValue": 100},{"roadName": "和平路", "roadSide": "", "dimmingValue": 0}]}"""
    init_prompt = """
    缺少道路名称或者调光值，请询问用户更多信息，引导用户输入道路名称和调光值
    """
    try:
        for roadModel in roadModelList:
            if roadModel.roadName is None or roadModel.roadName == "" or roadModel.dimmingValue is None or roadModel.dimmingValue == "":
                return init_prompt
            return "操作成功"
    except:
        return init_prompt
    return init_prompt

@tool
def pole_promote(lamppostModelList: List[LamppostModel]):
    """按灯杆编号进行开关灯。参数格式为json:{"lamppostModelList": [{"lamppostCode": "201", "dimmingValue": 100},{"lamppostCode": "202407051130", "dimmingValue": 100}]}"""
    init_prompt = """
        缺少灯杆编号或者调光值，请询问用户更多信息，引导用户输入灯杆编号和调光值
        """
    try:
        for lamppostModel in lamppostModelList:
            if lamppostModel.lamppostCode is None or lamppostModel.dimmingValue is None:
                return init_prompt
            return "操作失败"
    except:
        return init_prompt

@tool
def ask_for_more_info(question):
    """询问用户更多信息,对用户进行提问。参数question是字符串类型,格式：{"question": ".....？"} """
    #print(f"问题：{question}")
    return input("请输入更多信息：")

tools = [road_promote, pole_promote, ask_for_more_info]
rendered_tools = render_text_description(tools)
openai_tools = [convert_to_openai_tool(tool) for tool in tools]
#print(openai_tools)


tool_names = [i.name for i in tools]
parser = JsonOutputParser()
def tool_chain(model_output):
    tool_map = {tool.name: tool for tool in tools}
    chosen_tool = tool_map[model_output["name"]]
    return itemgetter("arguments") | chosen_tool

def get_args(res):
    args=res.split("\nAction:")[1].split("\nAction Input:")
    try:
        arguments = parser.parse(args[1].strip())
    except:
        arguments = {}
    return {"name": args[0].strip(), "arguments": arguments}

def prompt_react(input):
    PREFIX = f"""你是一个开关灯助手，只能回答开关灯相关的问题，如果用户提出的问题与开关灯无关，需要引导用户提问开关灯相关的问题。
    不要进行任何猜测和假设,如果需要用户回答问题,请将问题引导到ask_for_more_info
    不要进行任何猜测和假设,如果需要用户回答问题,请将问题引导到ask_for_more_info
    不要进行任何猜测和假设,如果需要用户回答问题,请将问题引导到ask_for_more_info
    
    不要有任何说明、解释、注释
    不要有任何说明、解释、注释
    不要有任何说明、解释、注释
    
    有以下工具可以调用：
        {openai_tools}
    请尽可能从用户输入的问题中找到合适的工具
    
    "Final Answer的作用只输出最终的答案,如果需要用户回答问题,请将问题引导到ask_for_more_info"
    "Final Answer的作用只输出最终的答案,如果需要用户回答问题,请将问题引导到ask_for_more_info"
    """
    prompt = f"""
        Use the following format:
        Question: 您必须回答的输入问题
        Thought: 你应该经常考虑该怎么做
        Action: 按顺序采取的行动应为以下操作之一：{tool_names}
        Action Input: 操作的输入，不要有任何解释和说明，直接输入，不能包含注释，不能包含注释，不能包含注释，不能包含注释，不能包含注释，不能有"//"
        Observation: 动作执行结果
        ... (this Thought/Action/Action Input/Observation can repeat N times)
        Thought: 我现在知道了最终的答案
        Final Answer: 原始输入问题的最终答案。需要用户输入信息时，不要输出这个最终答案
        
        Begin!
        
        Question: {input}
        """
    return [{"role": "system", "content": PREFIX},{"role": "user", "content": prompt}]


tool_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", """
            你是一个开关灯助手，只能回答开关灯相关的问题，如果用户提出的问题与开关灯无关，需要引导用户提问开关灯相关的问题。
            不要进行任何猜测和假设,如果需要用户回答问题,请将问题引导到ask_for_more_info
            不要进行任何猜测和假设,如果需要用户回答问题,请将问题引导到ask_for_more_info
            不要进行任何猜测和假设,如果需要用户回答问题,请将问题引导到ask_for_more_info
        """),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ]
)


#chain = prompt_react | tongyi_llm
query = input("请输入问题：")
p_r = prompt_react(query)

while True:
    agent = create_tool_calling_agent(local_llm, tools, prompt=tool_prompt)
    agent_executor = AgentExecutor(
    agent=agent, tools=tools, early_stopping_method="force", max_iterations=2,verbose=True
    )
    res = agent_executor.invoke({"input":query})
    print("\033[1;32;40m" + res["output"] + "\033[0m")
    query = input("请输入问题：")